from math import log2
import matplotlib.pyplot as plt

def shannonEnt(dataSet):
    num = len(dataSet)
    labelCount = {}
    for feature in dataSet:
        label = feature[-1]
        labelCount[label] = labelCount.get(label,0)+1
    shannonEnt = 0.0
    # print(labelCount)
    for key in labelCount:
        probability =  float(labelCount[key])/num
        shannonEnt -= probability*log2(probability)
    return shannonEnt

def createDataSet():
    dataSet = [[1,1,'yes'],[1,1,'yes'],[1,0,'no'],
                     [0,1,'no'],[0,1,'no']]
    labels = ['no surfacing','flippers']
    return dataSet,labels

def splitDataSet(dataSet,axis,value):
    retDataSet = []
    for feature in dataSet:
        feature = list(feature)
        if feature[axis] == value:
            reducedFeature = feature[:axis]
            reducedFeature.extend(feature[axis+1:])
            # a.extend(b) if b is list scattering it and adding into a in sequence
            retDataSet.append(reducedFeature)
    return retDataSet

def chooseBestFeatureToSplit(dataSet):
    numf = len(dataSet[0])-1
    baseEnt = shannonEnt(dataSet)
    bestInfoGain = 0.0
    bestF = -1
    for i in range(numf):
        featureList = [example[i] for example in dataSet]
        uniqueVals = set(featureList)
        newEnt = 0.0
        for value in uniqueVals:
            subDataSet = splitDataSet(dataSet,i,value)
            prob = len(subDataSet)/float(len(dataSet))
            newEnt += prob*shannonEnt(subDataSet)
        infoGain = baseEnt-newEnt
        if infoGain>bestInfoGain:
            bestInfoGain = infoGain
            bestF =i
    return bestF

def chooseBFTS(dataSet):
    numf = len(dataSet[0])-1
    baseEnt = shannonEnt(dataSet)
    bestF = -1
    for i in range(numf):
        featureList = [example[i] for example in dataSet]
        uniqueVals = set(featureList)
        newEnt = 0.0
        for value in uniqueVals:
            subDataSet = splitDataSet(dataSet,i,value)
            prob = len(subDataSet)/float(len(dataSet))
            newEnt += prob*shannonEnt(subDataSet)
        if newEnt<baseEnt:
            baseEnt = newEnt
            bestF = i
    return bestF

def majorityCnt(classList):
    classCount = {}
    for vote in classList:
        classCount[vote] = classCount.get(vote,0)+1
    sortedCount = sorted(classCount.items(),key=lambda x:x[1],reverse=True)
    return  sortedCount[0][0]

def createTree(dataSet,labels):
    classList = [example[-1] for example in dataSet]
    if classList.count(classList[0]) == len(classList):
        return classList[0]
    if len(dataSet[0])==1:
        return majorityCnt(classList)
    bestFeature = chooseBestFeatureToSplit(dataSet)
    bestFLable = labels[bestFeature]
    myTree = {bestFLable:{}}
    del(labels[bestFeature]) # python中字典和列表为地址引用，对起结构进行操作会影响函数外的值，而赋值则不会.赋值时用深拷贝
    featureVal = [example[bestFeature] for example in dataSet]
    uniqueVal = set(featureVal)
    for value in uniqueVal:
        myTree[bestFLable][value] = createTree(splitDataSet(dataSet,bestFeature,value),labels[:])
    return myTree


def getNumLeafs(myTree):
    num = 0
    firstStr = list(myTree.keys())[0]
    secondDict = myTree[firstStr]
    if type(secondDict).__name__=='dict':
        for key in secondDict.keys():
            if type(secondDict[key]).__name__=='dict':
                num += getNumLeafs(secondDict[key])
            else:
                num +=1
    return num

def getTreeDepth(myTree):
    mdepth = 0
    firstStr = list(myTree.keys())[0]
    secondDict  = myTree[firstStr]
    if type(secondDict).__name__=='dict':
        for key in secondDict.keys():
            if type(secondDict[key]).__name__=='dict':
                depth = 1+getTreeDepth(secondDict[key])
            else:
                depth = 1
            mdepth = max(mdepth,depth)
    return mdepth

def plotMidText(cntrPt,parentPt,txtString):
    xmid = (parentPt[0]-cntrPt[0])/2+cntrPt[0]
    ymid = (parentPt[1]-cntrPt[1])/2+cntrPt[1]
    plt.text(xmid,ymid,txtString)

def plotTree(myTree,parentPt,nodeTxt = ''):
    numLeafs = getNumLeafs(myTree)
    # depth = getTreeDepth(myTree)
    firstStr = list(myTree.keys())[0]
    cntrPt = (plotTree.xOff+(1+numLeafs)/2/plotTree.totalW,plotTree.yOff)
    plotMidText(cntrPt,parentPt,nodeTxt)
    params = {0: {'nodeTxt': firstStr, 'centerPt': cntrPt, 'parentPt': parentPt, 'nodeType': dict(boxstyle='sawtooth', fc='0.8')}}
    plt.annotate(params[0]['nodeTxt'], xy=params[0]['parentPt'], xycoords='axes fraction', xytext=params[0]['centerPt'], \
                 textcoords='axes fraction', va='center', ha='center', bbox=params[0]['nodeType'],
                 arrowprops=dict(arrowstyle='<-'))

    secondDict = myTree[firstStr]
    plotTree.yOff = plotTree.yOff-1/plotTree.totalD
    for key in secondDict.keys():
        if type(secondDict[key]).__name__=='dict':
            plotTree(secondDict[key],cntrPt,str(key))
        else:
            plotTree.xOff = plotTree.xOff+1/plotTree.totalW
            params = {0: {'nodeTxt': secondDict[key], 'centerPt': (plotTree.xOff,plotTree.yOff), 'parentPt': cntrPt,
                          'nodeType': dict(boxstyle='round4', fc='0.8')}}
            plt.annotate(params[0]['nodeTxt'], xy=params[0]['parentPt'], xycoords='axes fraction',
                         xytext=params[0]['centerPt'], \
                         textcoords='axes fraction', va='center', ha='center', bbox=params[0]['nodeType'],
                         arrowprops=dict(arrowstyle='<-'))
            plotMidText((plotTree.xOff,plotTree.yOff),cntrPt,str(key))
    plotTree.yOff = plotTree.yOff+1/plotTree.totalD


def createPlot(inTree):
    fig = plt.figure(1,facecolor='white')
    fig.clf()  # 清空绘图区
    plt.axis('off')    # 去掉坐标轴      去掉刻度 plt.xticks([])/plt.yticks([])
    # ax = plt.subplot(111,frameon=False,**dict(xticks=[],yticks=[]))
    plotTree.totalW = float(getNumLeafs(inTree))
    plotTree.totalD = float(getTreeDepth(inTree))
    plotTree.xOff = -0.5/plotTree.totalW
    plotTree.yOff = 1
    plotTree(inTree,(0.5,1.0))
    plt.show()


def classify(inputTree,labels,testVec):
    firstStr = list(inputTree.keys())[0]
    secondDict = inputTree[firstStr]
    index = labels.index(firstStr)
    classLabel = 'inputTree error'
    if type(secondDict).__name__ == 'dict':
        for key in secondDict.keys():
            if testVec[index] == key:
                if type(secondDict[key]).__name__ == 'dict':
                    classLabel = classify(secondDict[key],labels,testVec)
                else:
                    classLabel = secondDict[key]
    return classLabel









# -- examination --
'''

mydata, labels = createDataSet()
myTree = createTree(mydata,labels[:]) # 列表深拷贝
rs = classify(myTree,labels,[1,0])
print(rs)

myTree = {'no surfacing': {0: 'no', 1: {'flippers': {0: 'no', 1: {'flippers': {0: 'no', 1: 'yes'}}}},2:'maybe'}}
createPlot(myTree)




print(chooseBFTS(mydata))
print(getNumLeafs(createTree(mydata,labels)))



'''